
from pymongo import ASCENDING
from pev3.util.DataBase import DB_CONN
from pandas import DataFrame
from datetime import datetime

"""
数据处理子系统，主要完成的工作
1. 从外部数据源获取数据，并保存到数据库中
2. 为交易系统提供数据接口
"""

class DataModule:

    def __init__(self):
        self.db = DB_CONN

    """
    获取指定股票代码在固定周期的数据
    :param code: 股票代码
    :param index: 是否是指数
    :param autype: 复权类型，None - 不复权，qfq - 前复权, hfq - 后复权
    :param period: K线周期，D - 日线(默认值)，W - 周线， M - 月线，M1 - 1分钟，M5 - 5分钟
    :param startDate: 数据的开始日期
    :param endDate: 数据的结束日期
    :return: 包含K线的DataFrame
    """
    def getKData(self, code, index=None, autype=None, period='D', startDate=None, endDate=None):
        print("子系统：数据处理，操作：获取K线数据，状态：开始， 股票：%s, 级别: %s, 开始日期：%s, 结束日期: %s" %
              code, period, startDate, endDate)

        # 如果是指数，则从daily数据集中查询数据,如果不是指数，则根据复权类型从相应的数据集中查询数据
        if index:
            dailyCursor = self.db['daily'].find(
                {'code': code, 'index': True, 'date': {'$gte': startDate, '$lte': endDate}},
                sort=[{'date', ASCENDING}],
                projection={'_id': False},
                batch_size = 500
            )
        else:
            if autype is None:
                dailyCursor = self.db['daily'].find(
                    {'code': code, 'index': False, 'date': {'$gte': startDate, '$lte': endDate}},
                    sort=[{'date': ASCENDING}],
                    projection={'_id': False},
                    batch_size=500
                )
            else:
                dailyCursor = self.db['daily_' + autype].find(
                    {'code': code, 'index': False, 'date': {'$gte': startDate, '$lte': endDate}},
                    sort=[{'date': ASCENDING}],
                    projection={'_id': False},
                    batch_size=500
                )

        dailiesDF = DataFrame([daily for daily in dailyCursor])

        # index字段不返回
        dailiesDF.drop(['index'], 1, inplace=True)

        print("子系统：数据处理，操作：获取K线数据，状态：结束，股票：%s, 级别：%s, 开始日期：%s，结束日期： %s" %
              code, period, startDate, endDate)
        return dailiesDF

    """
    获取某一日所有股票的K线数据
    :param autype: 复权类型，None - 不复权，qfq - 前复权, hfq - 后复权
    :param period: K线周期，D - 日线(默认值)，W - 周线， M - 月线，M1 - 1分钟，M5 - 5分钟
    :param date: 日期
    :return: 包含K线的DataFrame
    """
    def getOneDayKData(self, autype=None, period='D', date=None):
        print("子系统：数据处理，操作：获取K线数据，状态：开始， 级别：%s, 日期：%s" %
              (period, date), flush=True)
        # 根据复权类型从相应的数据集中查询数据
        if autype is None:
            dailyCursor = self.db['daily'].find(
                {'index':False, 'date': date},
                projection={'_id': False},
                batch_size=500
            )
        else:
            dailyCursor = self.db['daily_' + autype].find(
                {'index': False, 'date': date},
                projection={'_id': False},
                batch_size=500
            )

        dailiesDF = DataFrame([daily for daily in dailyCursor])

        # index字段不返回
        dailiesDF.drop(['index'], 1, inplace=True)

        print("子系统：数据处理，操作：获取K线数据，状态：结束， 级别：%s, 日期：%s" %
              (period, date), flush=True)
        return dailiesDF

    """
    查询出所有股票在某个交易日的基本数据
    :param date: 指定的日期
    :return: 基本信息的DataFrame
    """
    def getStockBasic(self, date=None):

        if date is None:
            date = datetime.now().strftime('%Y-%m-%d')

        basicCursor = self.db['basic'].find(
            {'date': date},
            projection={'_id': False},
            batch_size=1000
        )

        basicDF = DataFrame([basic for basic in basicCursor])

        return basicDF

if __name__ == '__main__':
    print(DataModule().getStockBasic('2018-10-15'))
